package com.llmops.core.node;

import com.llmops.core.Context;
import com.llmops.core.Node;
import com.llmops.core.NodeProperty;
import com.llmops.core.listener.ChatModelLogListener;
import dev.langchain4j.model.input.structured.StructuredPrompt;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.structured.Description;
import dev.langchain4j.service.AiServices;
import lombok.Data;

import java.util.List;
import java.util.Map;

public class ClassificationNode extends Node {
    private ClassificationNodeProperty property = new ClassificationNodeProperty();
    @Data
    public static class ClassificationNodeProperty extends NodeProperty {
        private String apiUrl;
        private String apiKey;
        private String model;
    }

    static class ClassificationResult {

        @Description("分类的结果")
        private String classification;

        @Override
        public String toString() {
            return "[分类结果："+classification+"]";
        }
    }

    @StructuredPrompt("将用户输入:{{input}},分类为:{{classifications}}")
    class CreateRecipePrompt {
        private String input;
        private List<String> classifications;
    }

    interface Classification {
        ClassificationResult classification(CreateRecipePrompt prompt);
    }



    public ClassificationNode(String id,String title, Map<String, Object> config) {
        super(id, title,config);
        this.property.setApiUrl((String) config.getOrDefault("apiUrl", "https://api.deepseek.com/v1"));
        this.property.setApiKey((String) config.getOrDefault("apiKey","sk-0c3e6a9f9f7a4e63bb855290c544183c"));
        this.property.setModel((String) config.getOrDefault("model", "deepseek-chat"));
    }

    @Override
    public void execute(Context context) {
        OpenAiChatModel chatModel = OpenAiChatModel.builder()
                .baseUrl(this.property.getApiUrl())
                .apiKey(this.property.getApiKey())
                .modelName(this.property.getModel())
                .listeners(List.of(new ChatModelLogListener()))
                .build();


        Classification classification = AiServices.create(Classification.class, chatModel);

        CreateRecipePrompt prompt = new CreateRecipePrompt();
        prompt.input = context.getVariable(getPreviousNode().getId()+"_output");
        prompt.classifications =(List<String>) config.get("classifications");
        ClassificationResult anotherRecipe = classification.classification(prompt);
        context.setVariable(id+"_output",anotherRecipe.classification);
        logger.info(getId() + "_output={}", anotherRecipe.classification);
    }

    @Override
    public String getType() {
        return "classificationNode";
    }
}
